package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{ImputeEncoder, ImputeEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author: dongguosheng
  * @Date: 2019/3/22 15:00
  * @Review songhaicheng
  * @Email: guosheng_dong@sui.com
  */
class ImputeEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "a,b",
      | "outputCol": "output1,output",
      |	"features": "features",
      |	"strategy": "mean",
      | "missingValue": "NaN"
      |}
    """.stripMargin

  val imputer = new ImputeEncoder()
  val param: ImputeEncoderParam = imputer.parseParam(new ToolClass().encrypt(paramStr))

  test("ImputerEncoder parameter construction test") {
    assert(param.input_pt == null && param.output_pt == null &&
      "a,b".equals(param.inputCol) && "output1,output".equals(param.outputCol) &&
      "mean".equals(param.strategy) && param.missingValue.isNaN
    )
  }

  test("ImputerEncoder training function test") {
    val testingData = TestingDataGenerator.makeImputerData()
    val res = imputer.train(param, testingData).transform(testingData)
    assert(res.select("output").union(res.select("output1"))
      .rdd
      .map(_.getAs[Double](0))
      .collect()
      .count(_.isNaN) == 0
    )
  }
}
